from agent.agent import Agent
from env.env_def import UnitType, BLUE_AIRPORT_ID
from .cmd import Command
from .task import Task, TaskState
from .a2a import A2A
import numpy as np
from env.env_cmd import EnvCmd
import math

TAKEOFF_PATROL_POINT = [-50000, 0, 8000]
TYPE2STRING_MAP = {11: "AIR", 12: "AWCS", 13: "JAM", 15: "BOM"}


class A2ASupportTask(Task):
    def __init__(self, task_state):
        super().__init__(task_state)
        self.done_unit_ids = []

    def _attack(self, obs_blue):
        cmds = []
        unit_ids = []
        en_unit_ids = []
        for _, unit in self.units_map.items():
            for en_unit in obs_blue['qb']:
                if en_unit['LX'] in [11, 15]:
                    dist = unit.compute_2d_distance(en_unit['X'], en_unit['Y'])
                    if dist < 120000:
                        # cmds.append(Command.a2a_attack(unit.id, en_unit['ID']))
                        unit_ids.append(unit.id)
                        en_unit_ids.append(en_unit['ID'])
        # unit_ids = list(set(unit_ids))
        # en_unit_ids = list(set(en_unit_ids))
        # for unit_id in unit_ids:
        #     np.random.shuffle(en_unit_ids)
        #     for en_unit_id in en_unit_ids:
        #         cmds.append(Command.a2a_attack(unit_id, en_unit_id))
        return cmds


class A2ASupportPoint(A2ASupportTask):
    def __init__(self, max_unit_num, target_point, task_state=TaskState.PREPARE):
        super().__init__(task_state)
        self.max_unit_num = max_unit_num
        self.target_point = target_point

    def get_all_missiles(self):
        missile_num = 0
        for _, unit in self.units_map.items():
            missile_num += unit.get_missile_num()
        return missile_num

    def update_task_state(self):
        pass

    def update(self, alive_unit_map):
        self.update_units_map(alive_unit_map)

    def run(self, idle_unit_map, obs_blue):
        cmds = []
        attack_cmds = self._attack(obs_blue)
        cmds.extend(attack_cmds)
        if self.task_state != TaskState.FINISH:
            while len(self.units_map) < self.max_unit_num:
                if len(idle_unit_map) == 0:
                    break
                unit = idle_unit_map.popitem()[1]
                self.add_unit(unit)
            if len(attack_cmds) == 0:
                for unit_id, _ in self.units_map.items():
                    cmds.append(Command.area_patrol(unit_id, [self.target_point[0], self.target_point[1], 8000]))
        else:
            for unit_id, _ in self.units_map.items():
                cmds.append(Command.area_patrol(unit_id, TAKEOFF_PATROL_POINT))
            self.finish()
        return cmds


class BlueAgent(Agent):
    def __init__(self, name, config, **kwargs):
        super().__init__(name, config['side'])
        self.a2a_task = A2ASupportPoint(3, [0, 0])
        self.find = False

    def _get_waiting_aircraft_num(self, obs_blue, type_):
        if type_ not in TYPE2STRING_MAP.keys():
            return 0
        type_str = TYPE2STRING_MAP[type_]
        return obs_blue['airports'][0][type_str]

    """
        控制飞机起飞
    """

    def _take_off(self, obs_blue):
        cmds = []
        fly_types = [UnitType.AWACS, UnitType.A2G, UnitType.JAM]
        for type_ in fly_types:
            if self._get_waiting_aircraft_num(obs_blue, type_):
                cmds.append(
                    Command.takeoff_areapatrol(
                        BLUE_AIRPORT_ID, 1, type_, patrol_points=[-50000, -40000+20000*obs_blue["airports"][0]["AWCS"], 8000]))

        return cmds

    def _awacs_task(self, obs_blue):
        cmds = []
        patrol_points = [-180000, 150000, 8000]
        for unit in obs_blue['units']:
            if unit['LX'] == UnitType.AWACS:
                cmds.append(
                    Command.awacs_areapatrol(
                        unit['ID'], patrol_points))

                break
        return cmds

    def _get_units_map(self, obs_blue, type_):
        units_map = {}
        for info_map in obs_blue['units']:
            if info_map['LX'] == type_:
                units_map[info_map['ID']] = A2A(info_map)
        return units_map

    def step(self, sim_time, obs_blue, **kwargs):
        cmds = []
        # obs_red = obs_blue["red"]
        # obs_blue = obs_blue["blue"]
        for unit in obs_blue['units']:
            if unit['LX'] == UnitType.AWACS:
                cmds.append(
                    EnvCmd.make_awcs_radarcontrol(unit['ID'], 1))
        for unit in obs_blue['qb']:
            if(unit['ID']!=20001 and unit['ID']!=30001):
                self.find = True
                break
        if obs_blue["airports"][0]["AWCS"] > 0 and sim_time > 0:
            cmds.extend(self._take_off(obs_blue))
            # cmds.extend(self._awacs_task(obs_blue))  # 预警机巡逻
            return cmds

        # a2a_map = self._get_units_map(obs_blue, UnitType.A2A)
        # self.a2a_task.update(a2a_map)
        # for unit_id, _ in self.a2a_task.get_units_map().items():
        #     a2a_map.pop(unit_id)
        # cmds.extend(self.a2a_task.run(a2a_map, obs_blue))

        if self.find:
            # obs是观测的
            self_a2a = []
            enemy_a2a = []
            my_a2a = 0
            # 遍历蓝方的units，看是否是飞机
            for unit in obs_blue["units"]:
                if (unit["LX"] == UnitType.AWACS):
                    my_a2a += 1
                    self_a2a.append(unit)
            for unit in obs_blue["qb"]:
                if (unit["LX"] == UnitType.AWACS or unit["LX"] == 18):  # 18：未知空中目标
                    enemy_a2a.append(unit)
            for i in range(my_a2a):
                i_unit = self_a2a[i]
                for j_unit in enemy_a2a:
                    if self.get_distance(i_unit, j_unit) < 200000:# 开始逃跑距离
                        print("----==================================================================")
                        cmds.append(EnvCmd.make_areapatrol(i_unit["ID"], 2 * i_unit["X"] - j_unit["X"],
                                                           2 * i_unit["Y"] - j_unit["Y"], 2 * i_unit["Z"] - j_unit["Z"],
                                                           0, 1000, 1000, 800, 7200))
                        # 速度稍慢于红方，位置是往敌方飞机的反方向
        return cmds

    def reset(self):
        self.find = False
        pass

    def get_distance(self, unit1, unit2):
        d_x = unit1["X"] - unit2["X"]
        d_y = unit1["Y"] - unit2["Y"]
        return math.sqrt(d_x ** 2 + d_y ** 2)
